//////////////////////////////////////////////////////////////////
// MattiseOS Kernel
// Date: 2008
// Author: Matthew Iselin
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// at your option) any later version.
// 
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
// 
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.
//////////////////////////////////////////////////////////////////

// Includes
#include <mattise.h>
#include <stdint.h>
#include <string.h>
#include <mem/util.h>
#include <mem/heap.h>
#include <mem/page.h>
#include <exec/elf.h>
#include <kernio/io.h>
#include <process/common.h>

// Prototypes

// the sbrk system call - argv shouldn't be on the kernel heap
int sbrk( int incr, struct context* ctx );

// Functions

// entry routine for a userspace thread
void uspace_thread_entry( uint32_t neweip )
{
	// open the tty devices
	sys_open( "/dev/tty", 0 ); // read
	sys_open( "/dev/tty", 0 ); // write
	sys_open( "/dev/tty", 0 ); // err

#if DEBUG
	dprintf( "starting task %d, argc=%d, argv=%x\n", currprocess->id, currprocess->argc, currprocess->argv );
#endif

	// setup the current directory properly (default to the first hdd)
	strcpy( currprocess->currdir, "/hd0/" );

	uint32_t cr3 = 0;
	asm volatile("mov %%cr3,%0" : "=r" (cr3));

	// load an executable if needed
	if(currprocess->exe_name)
	{
		// load it in
		int32_t status = elf_load(currprocess->exe_name, &neweip);
		if(status != 0)
		{
			dprintf("[FATAL] loading binary %s failed: %d\n", currprocess->exe_name, status);

			// don't need the exe name anymore
			free(currprocess->exe_name);
			currprocess->exe_name = 0;

			kill(getpid());
			while(1);
		}

		// don't need the exe name anymore
		free(currprocess->exe_name);
		currprocess->exe_name = 0;
	}

	// check for a valid argv
	char** argv = (char**) currprocess->argv;
	if( argv )
	{
		// find out how many arguments were passed
		int argc;
		for( argc = 0; argv[argc] != 0; argc++ );

		// check that we do actually need to copy arguments
		if( currprocess->cr3 == GetKernelDirectory() )
		{
			// copying is a Bad Idea (tm) when working with kernel threads - no
			// seperate thread heap! so, setup the argv/argc in the context to
			// whatever was passed
			currprocess->argc = argc;
			currprocess->argv = (uint32_t) argv;
		}
		else
		{
			// allocate space on the new process heap
			char** newproc_argv = (char**) sbrk(sizeof(char*) * (argc + 1), currprocess);

			// allocate and copy arguments
			int _arg;
			for( _arg = 0; _arg < argc; _arg++ )
			{
				// allocate space first
				newproc_argv[_arg] = (char*) sbrk(strlen(argv[_arg]) + 1, currprocess);

				// and then copy the argument across
				memcpy( newproc_argv[_arg], argv[_arg], strlen( argv[_arg] ) + 1 );

				// free this argument
				free(argv[_arg]);
			}

			// the last argument is always 0
			newproc_argv[_arg] = 0;

			// setup the argv/argc in the context
			currprocess->argc = argc;
			currprocess->argv = (uint32_t) newproc_argv;

			// free the old argv from the heap
			free(argv);
		}
	}

	// valid EIP?
	if(neweip == 0)
		kill(getpid());
	
	// set the userstack
	uint32_t userstack_virt = GetChildUserStack(currprocess) + STACK_USER_VIRT;
	uint32_t userstack = kAllocateAndMap(SZ_PSTACK, userstack_virt, PF_PRESENT | PF_READWRITE | PF_USER);
	userstack += (SZ_PSTACK * 0x1000) - 4;
	currprocess->userstack = userstack;

	// set as a new region in the memory map
	if(currprocess->memmap == 0)
	{
		currprocess->memmap = (struct ctx_mem_map*) malloc(sizeof(struct ctx_mem_map));
		currprocess->memmap->next = 0;
	}
	else
	{
		struct ctx_mem_map* new_region = (struct ctx_mem_map*) malloc(sizeof(struct ctx_mem_map));
		new_region->next = currprocess->memmap;
		currprocess->memmap = new_region;
	}
	currprocess->memmap->base = userstack_virt;
	currprocess->memmap->type = MEMMAP_ELEM_STACK;
	currprocess->memmap->size = SZ_PSTACK * 0x1000;

	// setup the stack to pretend we're returning from an interrupt
	// ss:esp, eflags, cs:eip
	asm volatile( 	"push $0x23; \
					push %1; \
					push $0x3202; \
					push $0x1b; \
					push %0; \
					mov $0x23,%%ax; \
					mov %%ax,%%ds; \
					mov %%ax,%%es; \
					mov %%ax,%%fs; \
					mov %%ax,%%gs; \
					mov %2,%%edi; \
					mov %3,%%esi; \
					sti; \
					iret;"
					:
					:
					"r" (neweip),
					"r" (userstack),
					"r" (currprocess->argc),
					"r" (currprocess->argv)
					);

	// we shouldn't ever get here
	while( 1 );
}

// entry routine for virtual 8086 mode tasks
void vm86_thread_entry( uint32_t neweip, uint32_t userstack )
{
	/** vm86 stack frame on return:

	4 vm86 segment registers
	ss3:esp3
	eflags
	cs:eip

	**/

	// setup the stack to pretend we're returning from an interrupt
	printf( "vm86 task new eip: %x, new stack: %x\n", neweip, userstack );
	printf( "opcode there: %x\n", *((uint32_t*) neweip) );

	asm volatile( "push $0x0; push $0x0; push $0x0; push $0x0" ); // push VM86 cs/ds/es/fs
	asm volatile( "push $0x23; push %0" : : "r" (userstack) ); // push the user stack
	asm volatile( "push $0x20202" ); // push eflags (IF, VM and always set bits)
	asm volatile( "push $0x1b; push %0" : : "r" (neweip) ); // push the return EIP
	asm volatile( "mov $0x23,%ax; mov %ax,%ds; mov %ax,%es; mov %ax,%fs; mov %ax,%gs" ); // setup data segments
	asm volatile( "iret" );
	//printf( "without iretting...\n" );
	//while( 1 );
	/*
					push $0x10; \
					push %1; \*/

	asm volatile( 	"push $0x0; \
					push $0x0; \
					push $0x0; \
					push $0x0; \
					push $0x20002; \
					push $0x08; \
					push %0; \
					mov $0x10,%%ax; \
					mov %%ax,%%ds; \
					mov %%ax,%%es; \
					mov %%ax,%%fs; \
					mov %%ax,%%gs; \
					mov %2,%%edi; \
					mov %3,%%esi; \
					sti; \
					iret;"
					:
					:
					"r" (neweip),
					"r" (userstack),
					"r" (currprocess->argc),
					"r" (currprocess->argv)
					);

	// we shouldn't ever get here
	while( 1 );
}

// entry routine for a kernel thread
void kernel_thread_entry( uint32_t neweip )
{
	// open the tty devices
	if( currprocess->cr3 != GetKernelDirectory() )
	{
		int fd = sys_open( "/dev/tty", 0 ); // read
		fd = sys_open( "/dev/tty", 0 ); // write
		fd = sys_open( "/dev/tty", 0 ); // err
	}

#if DEBUG
	dprintf( "starting task %d, argc=%d, argv=%x\n", currprocess->id, currprocess->argc, currprocess->argv );
#endif

	// setup the current directory properly (default to the first hdd)
	strcpy( currprocess->currdir, "/hd0/" );

	// check for a valid argv
	char** argv = (char**) currprocess->argv;
	if( argv )
	{
		// find out how many arguments were passed
		int argc;
		for( argc = 0; argv[argc] != 0; argc++ );

		// check that we do actually need to copy arguments
		if( currprocess->cr3 == GetKernelDirectory() )
		{
			// copying is a Bad Idea (tm) when working with kernel threads - no
			// seperate thread heap! so, setup the argv/argc in the context to
			// whatever was passed
			currprocess->argc = argc;
			currprocess->argv = (uint32_t) argv;
		}
		else
		{
			// allocate space on the new process heap
			char** newproc_argv = (char**) sbrk(sizeof(char*) * (argc + 1), currprocess);

			// allocate and copy arguments
			int _arg;
			for( _arg = 0; _arg < argc; _arg++ )
			{
				// allocate space first
				newproc_argv[_arg] = (char*) sbrk(strlen(argv[_arg]) + 1, currprocess);

				// and then copy the argument across
				memcpy( newproc_argv[_arg], argv[_arg], strlen( argv[_arg] ) + 1 );

				// free this argument
				free(argv[_arg]);
			}

			// the last argument is always 0
			newproc_argv[_arg] = 0;

			// setup the argv/argc in the context
			currprocess->argc = argc;
			currprocess->argv = (uint32_t) newproc_argv;

			// free the argv
			free(argv);
		}
	}

	// are we loading an EXE?
	if(currprocess->exe_name)
	{
		// tell me
		int32_t status = elf_load(currprocess->exe_name, &neweip);
		if(status != 0)
			dprintf("[FATAL] loading binary %s failed: %d\n", currprocess->exe_name, status);

		// don't need the exe name anymore
		free(currprocess->exe_name);
		currprocess->exe_name = 0;

		// valid EIP?
		if(neweip == 0)
			kill(getpid());

		// run the process
		asm volatile(	"mov %1,%%esi; \
				mov %0,%%edi; \
				jmp *%2;"
				:
				:
				"r" (currprocess->argc),
				"r" (currprocess->argv),
				"r" (neweip)
				);
	}
	else
	{
		typedef void (*kern_thread_within_kernel)(int,char**);
		kern_thread_within_kernel f = (kern_thread_within_kernel) neweip;
		f( currprocess->argc, (char**) currprocess->argv );
	}

	// if the thread returns, die
	kill(getpid());
	while( 1 );
}

// the place where accidental exits from the above two functions go
void unsafe_return_from_thread_entry()
{
	kill( getpid() );
	while( 1 ) kSchedule(0);
}

// creates a new thread
int32_t kCreateThread( uint32_t parent, uint32_t startaddr, uint32_t flags, uint32_t prio, char** argv, struct pentry* parent_ctx, const char* exename )
{
#if DEBUG
	dprintf("%s\n", __FUNCTION__);
#endif

    // verify input
    if( ((int32_t) parent) == -1 )
        return -1;

    // verify that the process table is valid
    if( ptable == 0 )
        return -1;

    // grab the starting eflags (so we can restore interrupts properly)
    uint32_t __START_EFLAGS = 0;
    asm volatile( "pushf; popl %0" : "=r" (__START_EFLAGS) );

    // we can't be interrupted here
    asm volatile( "cli" );

    // iterate to the end of the process list
    struct pentry* ent = 0, *prev = 0, *parentptr = 0;
	if(parent_ctx == 0)
	{
		int32_t id = 0;
		for( ent = ptable; ent != 0; prev = ent, ent = ent->next, id++ )
		{
			if( ent && ent->me && ent->me->id == parent && parentptr == 0 )
			{
				parentptr = ent;
				break;
			}
		}
	}
	else
		parentptr = parent_ctx;

#if DEBUG
	dprintf( "got prev=%x, ent=%x, parentptr=%x\n", prev, ent, parentptr );
#endif

	// check that the parent pointer is correct
    if( parentptr == 0 )
        return -1;

    // now we have a valid pointer to the end of the list in prev,
    // create a new context
    struct context* newcntx = (struct context*) malloc( sizeof( struct context ) );

    // check for a valid new context
    if( newcntx == 0 )
        return -1;

    // get a new process table entry
    struct pentry* newentry = (struct pentry*) malloc( sizeof( struct pentry ) );

    // check for a valid entry
    if( newentry == 0 )
    {
        free( newcntx );
        return -1;
    }

    // set up the entry
#if DEBUG
	dprintf( "linking into list\n" );
#endif

	if(p_readyq == 0)
	{
		p_readyq = newentry;
		p_readyq->prev = p_readyq->next = 0;
	}
	else
	{
		p_readyq->prev = newentry;
		newentry->prev = 0;
		newentry->next = p_readyq;
		p_readyq = newentry;
	}
	newentry->me = newcntx;

#if DEBUG
	dprintf( "setting up entry [newcntx=%x]\n", newcntx );
#endif

    // set up the new context
    newcntx->id = lastpid++;
    newcntx->isprocess = 0;

#if DEBUG
	dprintf( "getting cr3 [parentptr->me=%x]\n", parentptr->me );
#endif

	newcntx->ent = newentry;
    newcntx->cr3 = parentptr->me->cr3;
	newcntx->parent = parentptr->me;
    newcntx->parentid = parentptr->me->id;
    newcntx->childid = -1;
    newcntx->firstrun = 1;
    newcntx->numchildren = 0;
    newcntx->priority = prio;
	newcntx->wake_on_die = 0;

    // setup the sleep stuff
    newcntx->sleep_start_time = 0;
    newcntx->sleep_end_time = 0;
    newcntx->sleep_curr_time = 0;

#if DEBUG
	dprintf( "setting up fds\n" );
#endif

    // setup the file descriptor table
	size_t fdi;
	for( fdi = 0; fdi < 512; fdi++ )
		newcntx->filedescs[fdi].avail = 1;

#if DEBUG
	dprintf( "done basic setup\n" );
#endif

    // do we need to continue?
    if( flags & TF_CTXONLY )
	{
		// new child much?
		parentptr->me->numchildren++;

		// add the pid to the parent list
		parentptr->me->childpids = (uint32_t*) realloc( parentptr->me->childpids, (parentptr->me->numchildren) * sizeof( uint32_t ) );
		parentptr->me->childpids[parentptr->me->numchildren-1] = newcntx->id;

		// we can be interrupted again, but only enable interrupts if they were on before
		if( __START_EFLAGS & 0x200 )
			asm volatile( "sti" );

		// context made, leave it at that
		return newcntx->id;
	}

    // setup the heap
    if( flags & TF_USPACE || parentptr->me->cr3 != GetKernelDirectory() )
    {
    	newcntx->heap_base = 0x10000000; // at the 256 MB mark
		// 0x800000; // after the kernel 0-8 MB section
    	newcntx->heap_curr = 0; // no heap space allocated yet
    }

	// put the argv into the context, the thread entry will handle that
	newcntx->argv = (uint32_t) argv;

	///////////////////////////////////////////// ACTUAL PROCESS INIT /////////////////////////////////////////////

    // create a stack - make it a page long (remember this is the BOTTOM, you add 0x1000 for the stack pointer)
	uint32_t* newstack;
	uint32_t stack_size, kstack_size;
	if( newcntx->cr3 == GetKernelDirectory() )
	{
		kstack_size = stack_size = SZ_KSTACK; // kernel threads can't have too much stack space, but they'll choke if they don't have enough
	}
	else if( flags & TF_VM86 )
		stack_size = 2; // 8k is quite a bit of space for realmode programs
	else
	{
		stack_size = SZ_PSTACK;
		kstack_size = SZ_KSTACK;
	}

    // setup the context boolean stuff
    if( flags & TF_VM86 )
        newcntx->vm86task = 1;
    else
        newcntx->vm86task = 0;

	// memory map:
	// kernel: 0 -> 12 MB
	// heap: 0x1000000 -> 0x50000000
	// kernel heap stack: 0x50000000 -> 0x90000000
	// paging code temporary map positions: 0x90000000 -> 0xA0000000
	// drivers: 0xA0000000 -> 0xB0000000
	// kernel heap: 0xB0000000 -> 0xC0000000
	// executable image: 0xC0000000 -> 0xD0000000
	// user stacks: 0xD0000000 -> 0xE0000000
	// UNUSED SPACE: 0xE0000000 -> 0xF0000000
	// page tables: 0xFFC00000 -> 0xFFFFE000
	// page directory: 0xFFFFF000

    // calculate virtual addresses for everything
	//int numChildren = parentptr->me->numchildren - 1;
    uint32_t mainstack_virt = 0;
    uint32_t userstack_virt = 0;
    if( flags & TF_VM86 ) // virtual 8086 mode tasks need to think that they're under 1 MB ;)
        mainstack_virt = STACK_KERN_VIRT; //0x6000; // remember though, this is the bottom
    if( flags & TF_USPACE && flags & TF_VM86 )
        userstack_virt = 0x8000;
    if( ! ( flags & TF_VM86 ) ) // only special case is when we have a virtual 8086 mode task
        mainstack_virt = GetChildKernelStack(newcntx);

	// save the stack position
	newcntx->kernstack = mainstack_virt;

	// calculate the kernel stack for this task
	newstack = (uint32_t*) mainstack_virt; // kAllocateAndMapPageDirTwice(0xfffff000, newcntx->cr3, kstack_size, 0xf0000000, mainstack_virt, PF_PRESENT | PF_READWRITE | PF_USER, PF_PRESENT | PF_READWRITE | PF_USER);

    // put the stack into a known address
    uint32_t newstack_bottom = (uint32_t) newstack;

	// allocate space - this is actually the kernel stack, the userland
	// stack is allocated later
	newcntx->kernstack_phys = newstack_bottom;
	newcntx->kernstack_sz = kstack_size;

#if DEBUG
	dprintf( "Virtual main stack at %x\n", mainstack_virt + (0x1000 * (parentptr->me->numchildren - 1)) );
#endif

    // now that the kernel stack is setup, set it up properly
	newstack = (uint32_t*) (newstack_bottom + (0x1000 * kstack_size));
#if DEBUG
	dprintf("newstack = %x\n", newstack);
#endif

    // code segment to use
    uint32_t threadcs = 0;
    if( flags & TF_USPACE )
        threadcs = 0x18 | 3; // ring3
    else
        threadcs = 0x08;
    if( flags & TF_VM86 )
        threadcs = 0;

	// check to see if we need to modify the start address
	if( flags & TF_VM86 )
	{
		// TODO: write VM86 stuff
	}

	// the thread entry routines use this address to jump to
	newstack--;
	*newstack = startaddr;

	// this must be pushed as the return EIP from the entry routines
	// it basically loops forever or kills the thread, which is safer
	newstack--;
	*newstack = (uint32_t) unsafe_return_from_thread_entry;

    // setup the EIP
	if( flags & TF_VM86 )
		newcntx->eip = (uint32_t) vm86_thread_entry;
    else if( flags & TF_USPACE )
		newcntx->eip = (uint32_t) uspace_thread_entry;
	else
		newcntx->eip = (uint32_t) kernel_thread_entry;

    // figure out the difference between the two locations
    uint32_t stackdiff = ((uint32_t) newstack) - newstack_bottom;

    // set the stack location in the context data
	newcntx->esp = mainstack_virt + stackdiff;
    newcntx->ebp = (uint32_t) newcntx->esp;

	// we're ready to go now
    newcntx->status = PSTATE_READY;

    // new child much?
    parentptr->me->numchildren++;

	// add the pid to the parent list
	parentptr->me->childpids = (uint32_t*) realloc( parentptr->me->childpids, (parentptr->me->numchildren) * sizeof( uint32_t ) );
	parentptr->me->childpids[parentptr->me->numchildren - 1] = newcntx->id;

	// copy and set the exectuable name if needed
	if(exename)
	{
		newcntx->exe_name = (char*) malloc( strlen( exename ) + 1 );
		strcpy( newcntx->exe_name, exename );
	}
	else
		newcntx->exe_name = 0;

    // we can be interrupted again, but only enable interrupts if they were on before
    if( __START_EFLAGS & 0x200 )
        asm volatile( "sti" );

#if DEBUG
    dprintf( "new thread context pointer %x, esp is at %x, parent nc is %d\n", newcntx, (uint32_t) newcntx->esp, parentptr->me->numchildren );
#endif

    // return our id
    return newcntx->id;
}
